
#include "tftp-exception.hpp"
#include "tftp-packet.hpp"
#include "sock-inet.hpp"

#include <netinet/in.h>
#include <string.h>

////////////////////////////////////////////////////////////////////////////////

tftp_packet::tftp_packet(const uint8_t* packet, size_t len)
    : packet(packet), len(len)
{
    if(len < 4)
    {
        throw tftp_exception_malformed();
    }
}

const uint8_t* tftp_packet::get_packet() const
{
    return packet;
}

size_t tftp_packet::get_len() const
{
    return len;
}

uint16_t tftp_packet::get_opcode() const
{
    return ntohs(((uint16_t*)packet)[0]);
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_flow::tftp_packet_flow(const uint8_t* packet, size_t len, uint16_t expected_opcode, const char* invalid_opcode_msg)
    : tftp_packet(packet, len)
{
    if(get_opcode() != expected_opcode)
    {
        throw tftp_exception_opcode(this, invalid_opcode_msg);
    }
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_request::tftp_packet_request(const uint8_t* packet, size_t len, uint16_t expected_opcode, const char* invalid_opcode_msg)
    : tftp_packet_flow(packet, len, expected_opcode, invalid_opcode_msg)
{
    size_t filename_len;

    if((filename_len = strnlen((filename = (char*)packet+2), len-2)) == len-2 ||
        strnlen((mode = (char*)filename+filename_len+1), len-2-filename_len-1) == len-2-filename_len-1)
    {
        throw tftp_exception_malformed();
    }
    else if(strcmp(mode, "octet") != 0)
    {
        throw tftp_exception_mode();
    }
}

const char* tftp_packet_request::get_filename() const
{
    return filename;
}

const char* tftp_packet_request::get_mode() const
{
    return mode;
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_read::tftp_packet_read(const uint8_t* packet, size_t len)
    : tftp_packet_request(packet, len, 1, "invalid read packet")
{
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_write::tftp_packet_write(const uint8_t* packet, size_t len)
    : tftp_packet_request(packet, len, 2, "invalid write packet")
{
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_transmission::tftp_packet_transmission(const uint8_t* packet, size_t len, uint16_t expected_opcode, const char* invalid_opcode_msg, uint16_t expected_sn)
    : tftp_packet_flow(packet, len, expected_opcode, invalid_opcode_msg)
{
    if(get_sn() != expected_sn)
    {
        throw tftp_exception_sn();
    }
}

////////////////////////////////////////////////////////////////////////////////

uint16_t tftp_packet_transmission::get_sn() const
{
    return ntohs(((uint16_t*)packet)[1]);
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_ack::tftp_packet_ack(const uint8_t* packet, size_t len, uint16_t expected_sn)
    : tftp_packet_transmission(packet, len, 4, "invalid ack packet", expected_sn)
{
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_data::tftp_packet_data(const uint8_t* packet, size_t len, uint16_t expected_sn)
    : tftp_packet_transmission(packet, len, 3, "invalid data packet", expected_sn)
{
}

const uint8_t* tftp_packet_data::get_data() const
{
    return packet + 4;
}

size_t tftp_packet_data::get_datalen() const
{
    return len - 4;
}

std::ofstream& operator<<(std::ofstream& stream, const tftp_packet_data& p)
{
    stream.write((char*)p.get_data(), p.get_datalen());

    if(stream.fail())
    {
        switch(errno)
        {
        case ENOSPC:
            throw exception_disk_full();

        default:
            throw exception_unknown();
        }
    }

    return stream;
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_error::tftp_packet_error(const uint8_t* packet, size_t len)
    : tftp_packet_flow(packet, len, 5, "invalid error packet")
{
    if(strnlen((msg = (char*)packet+4), len-4) == len-4)
    {
        throw tftp_exception_malformed();
    }
}

uint16_t tftp_packet_error::get_code() const
{
    return ntohs(((uint16_t*)packet)[1]);
}

const char* tftp_packet_error::get_msg() const
{
    return msg;
}

////////////////////////////////////////////////////////////////////////////////

tftp_packet_factory::tftp_packet_factory()
    : cursor(0), addr(0, 0)
{
}

tftp_packet_factory::tftp_packet_factory(const sock_inet::address& addr)
    : cursor(0), addr(addr)
{
}

const uint8_t* tftp_packet_factory::get_packet() const
{
    return buffer;
}

uint8_t* tftp_packet_factory::get_packet()
{
    return buffer;
}

size_t tftp_packet_factory::get_len() const
{
    return cursor;
}

const sock_inet::address& tftp_packet_factory::get_address() const
{
    return addr;
}

tftp_packet_factory& tftp_packet_factory::operator<<(uint16_t x)
{
    if(cursor+sizeof(x) > sizeof(buffer))
    {
        throw std::runtime_error("buffer overrun");
    }

    *(uint16_t*)(buffer+cursor) = htons(x);
    cursor += sizeof(x);

    return *this;
}

tftp_packet_factory& tftp_packet_factory::operator<<(const char* x)
{
    size_t len = strlen(x)+1;

    if(cursor+len > sizeof(buffer))
    {
        throw std::runtime_error("buffer overrun");
    }

    memcpy(buffer+cursor, x, len);
    cursor += len;

    return *this;
}

tftp_packet_factory& tftp_packet_factory::operator<<(std::ifstream& stream)
{
    stream.read((char*)(buffer+cursor), (std::streamsize)(sizeof(buffer)-cursor));
    cursor += stream.gcount();

    return *this;
}

sock_inet_udp& operator<<(sock_inet_udp& s, const tftp_packet_factory& p)
{
    return s << sock_inet_udp::packet(p.get_packet(), p.get_len(), p.get_address());
}

////////////////////////////////////////////////////////////////////////////////
